from abc import ABC, abstractmethod

from torch.utils.data import DataLoader
from tqdm import tqdm

from ..runner.runner import Runner
from ..utils.utils import to_device


class Worker(ABC):
    DESC_TEMPLATE = "{job_name} for epoch [{step}]"

    def __init__(self, dataloader: DataLoader) -> None:
        super().__init__()
        self.dataloader = dataloader

    def set_runner(self, runner: Runner):
        self.runner = runner

    def work(self):
        self.before_work_hook()
        for batch_data in tqdm(self.dataloader, desc=self.get_desc()):
            batch_data = to_device(batch_data, self.runner.device)
            self.process_batch(batch_data)
        self.after_work_hook()

    @abstractmethod
    def process_batch(self, batch_data):
        raise NotImplementedError

    @abstractmethod
    def before_work_hook(self) -> None:
        pass

    @abstractmethod
    def after_work_hook(self) -> None:
        pass

    @abstractmethod
    def get_desc(self) -> str:
        """
        进度条中的任务说明字符串
        """
        raise NotImplementedError
